from typing import Union, Dict, Any
from fastapi import APIRouter, HTTPException, BackgroundTasks
from app.services.aeds_api_service import AnomalousEdgeDetectionService
from app.schemas.dhgam.aeds_schema import (
    AEDSRequest, 
    AEDSResponse, 
    AEDSMiddleResponse, 
    AnomalousEdge,
    GraphOutput
)
from app.schemas.facenet_base_schema import TaskStatus
from app.tools.redis_client import redis_client
import logging
import networkx as nx

# 配置日志
logger = logging.getLogger(__name__)
router = APIRouter()
service = AnomalousEdgeDetectionService()

@router.post("/create", response_model=AEDSMiddleResponse)
async def create_task(request: AEDSRequest) -> AEDSMiddleResponse:
    """
    创建新的异常边检测任务并存储在Redis中
    
    Args:
        request (AEDSRequest): 包含任务ID、回调URL和输入参数的请求对象
        
    Returns:
        AEDSMiddleResponse: 包含任务初始状态的响应对象
        
    Raises:
        HTTPException: 当任务创建失败时抛出500错误
    """
    try:
        task_response = AEDSMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params,
            metrics=[]
        )
        await redis_client.set_data(request.task_id, task_response.dict())
        return task_response
    except Exception as e:
        logger.error(f"创建异常边检测任务失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"创建任务失败: {str(e)}")

async def add_task_to_background(
    background_tasks: BackgroundTasks,
    input_params: Dict[str, Any],
    task_id: str
) -> None:
    """
    将异常边检测任务添加到后台执行队列
    
    Args:
        background_tasks: FastAPI的后台任务对象
        input_params: 算法输入参数
        task_id: 任务唯一标识
    """
    background_tasks.add_task(run_task_in_background, input_params, task_id)

@router.post("/execute/{task_id}", response_model=AEDSMiddleResponse)
async def execute_task(
    task_id: str, 
    background_tasks: BackgroundTasks
) -> AEDSMiddleResponse:
    """
    执行指定ID的异常边检测任务
    
    Args:
        task_id: 任务唯一标识
        background_tasks: FastAPI的后台任务对象
        
    Returns:
        AEDSMiddleResponse: 任务执行状态的响应对象
        
    Raises:
        HTTPException: 当任务不存在时抛出404错误，执行失败时抛出500错误
    """
    try:
        task_data = await redis_client.get_data(task_id)
        if not task_data:
            raise HTTPException(status_code=404, detail="任务不存在")

        input_params = task_data.get("input_params")
        result_key = f'{task_id}_result'
        
        # 更新任务状态为执行中
        running_status = {
            "task_id": task_id,
            "task_callback_url": task_data.get("task_callback_url"),
            "task_status": TaskStatus.RUNNING,
            "task_progress": 0,
            "input_params": input_params,
            "metrics": []
        }
        await redis_client.set_data(result_key, running_status)
        
        # 添加到后台任务队列
        await add_task_to_background(background_tasks, input_params, task_id)
        
        return AEDSMiddleResponse(**running_status)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"执行异常边检测任务失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"执行任务失败: {str(e)}")

@router.get("/result/{task_id}", response_model=Union[AEDSResponse, AEDSMiddleResponse])
async def get_result(task_id: str) -> Union[AEDSResponse, AEDSMiddleResponse]:
    """
    获取异常边检测任务执行结果
    
    Args:
        task_id: 任务唯一标识
        
    Returns:
        Union[AEDSResponse, AEDSMiddleResponse]: 
            - 如果任务完成，返回AEDSResponse
            - 如果任务进行中，返回AEDSMiddleResponse
            
    Raises:
        HTTPException: 当结果不存在时抛出404错误，获取失败时抛出500错误
    """
    try:
        result_key = f'{task_id}_result'
        result_data = await redis_client.get_data(result_key)
        if not result_data:
            raise HTTPException(status_code=404, detail="结果不存在")
        
        return (AEDSResponse(**result_data) 
                if "output_params" in result_data 
                else AEDSMiddleResponse(**result_data))
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取异常边检测结果失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"获取结果失败: {str(e)}")

async def run_task_in_background(input_params: Dict[str, Any], task_id: str) -> None:
    """
    在后台执行异常边检测算法任务
    
    Args:
        input_params: 算法输入参数
        task_id: 任务唯一标识
    """
    try:
        # 更新进度
        await update_progress(task_id, 10, "正在转换知识图谱为NetworkX格式")
        
        # 将知识图谱转换为NetworkX格式
        graph_data = input_params.get("graph")
        threshold = input_params.get("threshold", 0.7)
        has_labels = input_params.get("has_labels", False)
        
        G = nx.Graph()
        
        # 添加节点
        for node in graph_data.get("nodes", []):
            node_attrs = {"name": node["name"]}
            if node.get("properties"):
                for prop in node["properties"]:
                    node_attrs[prop["name"]] = prop["value"]
            G.add_node(node["id"], **node_attrs)
        
        # 添加边
        for edge in graph_data.get("edges", []):
            edge_attrs = {"name": edge["name"]}
            if edge.get("properties"):
                for prop in edge["properties"]:
                    edge_attrs[prop["name"]] = prop["value"]
            G.add_edge(edge["from_node"], edge["to_node"], **edge_attrs)
        
        await update_progress(task_id, 30, "图转换完成，开始训练模型")
        
        # 如果图中的边已有标签，则进行训练
        if has_labels:
            service.fit(G)
        
        await update_progress(task_id, 60, "模型训练完成，开始检测异常边")
        
        # 检测异常边
        anomalous_edges_dict = service.detect(G, threshold)
        
        # 转换结果格式
        anomalous_edges = []
        for (u, v), score in anomalous_edges_dict.items():
            anomalous_edges.append(AnomalousEdge(
                from_node=u,
                to_node=v,
                anomaly_score=score
            ))
        
        await update_progress(task_id, 90, "异常边检测完成，正在生成结果")
        
        # 准备输出数据
        output_params = GraphOutput(
            anomalous_edges=anomalous_edges,
            total_edges=G.number_of_edges(),
            anomalous_count=len(anomalous_edges)
        )
        
        # 创建最终结果
        result = {
            "task_id": task_id,
            "task_status": TaskStatus.COMPLETED,
            "task_progress": 100,
            "output_params": output_params.dict()
        }
        
        await redis_client.update_data(f'{task_id}_result', result)
    except Exception as e:
        logger.error(f"异常边检测后台任务执行失败: {str(e)}", exc_info=True)
        # 更新任务状态为失败
        error_status = {
            "task_id": task_id,
            "task_status": TaskStatus.FAILED,
            "error_message": str(e)
        }
        await redis_client.update_data(f'{task_id}_result', error_status)

async def update_progress(task_id: str, progress: int, message: str) -> None:
    """更新任务进度
    
    Args:
        task_id: 任务ID
        progress: 完成百分比
        message: 进度消息
    """
    result_key = f'{task_id}_result'
    result_data = await redis_client.get_data(result_key)
    if result_data:
        result_data["task_progress"] = progress
        result_data["progress_message"] = message
        await redis_client.set_data(result_key, result_data) 